import sys
from esidlm.learner.dsnet import DSNETLearner
sys.path.append("../")

def dsnet_train():
    DSNET_TRAINING_CONFIG = {

        "global": {
            "seed": 42,
            "output_folder": "densenet",
        },

        "data": {
            "train_data": "../data/Himawari8_2023_2024_samplebased_train.csv",
            "valid_data": "../data/Himawari8_2023_2024_samplebased_valid.csv",
            "test_data": "../data/Himawari8_2023_2024_samplebased_test.csv",
            "cont_cols": ['band_1', 'band_2', 'band_3', 'band_4', 'band_5', 'band_6', 'band_7',
                          'band_8', 'band_9', 'band_10', 'band_11', 'band_12', 'band_13', 'band_14',
                          'band_15', 'band_16', 'band_17', 'band_18', 'band_19', 'band_20',
                          'band_17_sin', 'band_18_sin', 'band_19_sin', 'band_20_sin',
                          'band_17_cos', 'band_18_cos', 'band_19_cos', 'band_20_cos',
                          'band_9_13', 'band_9_15', 'band_10_15', 'band_11_13', 'band_13_15'],
            "cate_cols": ['hour_loc'],
            "target_col": "PM25",
        },

        "dataloader": {
            "batch_size": 512,
            "num_workers": 4,
        },

        "model": {
            "net": {
                "d_embed": 4,
                "d_model": 2,
                "n_layers": 2,
                "p_drop": 0.1,
                "act_fn": "relu"
            },
            "optimizer": {
                "lr": 0.0003,
                "weight_decay": 1e-4,
            },
        },

        "callback": {
            "model_checkpoint": {
                "save_top_k": 1,
                "monitor": "valid_r2",
                "mode": "max",
                "verbose": True
            },
            "early_stopping": {
                "monitor": "valid_r2",
                "mode": "max",
                "patience": 800,
                "verbose": True
            }
        },

        "trainer": {
            "max_epochs": 800,
            "accelerator": "cpu",
            "devices": 1,
            "deterministic": True
        }
    }
    learner = DSNETLearner(DSNET_TRAINING_CONFIG)
    learner.run_model_training()

def dsnet_inference():
    DSNET_INFERENCE_CONFIG = {
        "tif_config": {
            'cloud_mask': 1,
            'csv2tif': 1,
        },
        "global": {
            "seed": 42,
            "output_folder": "/mnt/d/g/output2",
        },

        "data": {
            "inference_folder": "../data/Himawari8_2023_2024_samplebased_test.csv",
            "cont_cols": ['band_1', 'band_2', 'band_3', 'band_4', 'band_5', 'band_6', 'band_7',
                          'band_8', 'band_9', 'band_10', 'band_11', 'band_12', 'band_13', 'band_14',
                          'band_15', 'band_16', 'band_17', 'band_18', 'band_19', 'band_20',
                          'band_17_sin', 'band_18_sin', 'band_19_sin', 'band_20_sin',
                          'band_17_cos', 'band_18_cos', 'band_19_cos', 'band_20_cos',
                          'band_9_13', 'band_9_15', 'band_10_15', 'band_11_13', 'band_13_15'],
            "cate_cols": ["hour_loc"],
            "target_col": "PM25",
        },

        "dataloader": {
            "batch_size": 1024,
            "num_workers": 4,
        },

        "model": {
            "model_checkpoint_path": "/mnt/d/g/output2/lightning_logs/version_11/checkpoints/epoch=79-step=18720.ckpt"
        },

        "trainer": {
            "accelerator": "gpu",
            "devices": 1,
        }
    }
    learner = DSNETLearner(DSNET_INFERENCE_CONFIG)
    learner.run_model_inference()


if __name__ == "__main__":
    dsnet_train()
